from tools.search_tool import get_UID
from langchain_ollama import ChatOllama
from langchain_classic.agents import AgentType, initialize_agent, Tool
from langchain_core.prompts import PromptTemplate

def lookup_V(flower_type: str):
    llm = ChatOllama(
        model="deepseek-r1:8b",
        base_url="http://localhost:11434",
        temperature=0.5,
        num_predict=512
    )

    template = """given the {flower} I want you to get a related 微博 UID. 
    You will receive search results in JSON format with 'url' fields. 
    Look for URLs starting with 'https://weibo.com/u/'. 
    Extract the number after '/u/' - that's the UID. 
    Your answer should contain only the UID number, no other text."""
    # 完整的提示模板
    prompt_template = PromptTemplate(input_variables=["flower"], template=template)

    # 代理的工具
    tools = [Tool(
        name="Search for Weibo UID",
        func=get_UID,
        description="useful for when you need to find Weibo UID for a specific flower",
    )]

    agent = initialize_agent(
        tools=tools,
        llm=llm,
        agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
        verbose=True,
        handle_parsing_errors=True,
        max_iterations=4,
        early_stopping_method="generate"
    )

    ID = agent.invoke(prompt_template.format_prompt(flower=flower_type))
    print('ID', ID)
    
    return ID['output']